package com.yunban.project.api.text_check.service.impl;

import com.alibaba.fastjson2.JSON;
import com.yunban.project.api.text_check.domain.CorrectionConfig;
import com.yunban.project.api.text_check.dto.GradioAiResp;
import com.yunban.project.api.text_check.dto.coze.CozeChatResp;
import com.yunban.project.api.text_check.dto.coze.CozeChatResult;
import com.yunban.project.api.text_check.dto.coze.CozeUnifiedResp;
import com.yunban.project.api.text_check.service.AiCorrectionService;
import com.yunban.project.api.text_check.utils.CorrectUtils;
import jakarta.annotation.Resource;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import net.lab1024.sa.base.common.exception.BusinessException;
import org.springframework.http.*;
import org.springframework.stereotype.Service;
import org.springframework.util.ObjectUtils;
import org.springframework.util.StringUtils;
import org.springframework.web.client.RestTemplate;

import java.util.ArrayList;
import java.util.List;

/**
 * @author: lijiajun
 * @date: 2024-11-25
 * @version: 1.0
 */
@Slf4j
@Service
public class CozeAiCorrectionService implements AiCorrectionService {

    @Resource
    private RestTemplate restTemplate;

    private static final String GET_CHAT_DETAILS_URL = "https://api.coze.cn/v3/chat/retrieve?conversation_id=%s&chat_id=%s";
    private static final String GET_CHAT_RESULT_URL = "https://api.coze.cn/v3/chat/message/list?conversation_id=%s&chat_id=%s";

    @Override
    public String correction(String text, CorrectionConfig config, String taskId) {
        CorrectUtils.printStartTime(config.getCode());
        String sourceText = text;
        // 构建提示词
        text = config.getPromptPrefix() + text + config.getPromptSuffix();
        String secretToken = config.getApiSecret();
        String url = config.getUrlPath();
        String botId = config.getBotId();
        if (!StringUtils.hasText(url) || !StringUtils.hasText(secretToken) || !StringUtils.hasText(botId)) {
            throw new BusinessException(config.getName() + "Coze 配置信息有误，请联系管理员~");
        }
        // 构建Coze请求参数
        CozeRequestParam params = new CozeRequestParam(text, botId);
        try {
            // header中 增加bearer token
            HttpHeaders headers = new HttpHeaders();
            headers.setContentType(MediaType.APPLICATION_JSON);
            headers.set("Authorization", "Bearer " + secretToken);
            // 创建 HttpEntity 对象，包含请求体和请求头
            HttpEntity<String> req = new HttpEntity<>(JSON.toJSONString(params), headers);
            CozeUnifiedResp result = this.restTemplate.postForObject(url, req, CozeUnifiedResp.class);
            log.info("coze text correct chatId:{}", result);
            if (ObjectUtils.isEmpty(result) || result.getCode() != 0) {
                log.error("coze text correction result error:{}", JSON.toJSONString(result));
                return null;
            }
            CozeChatResp cozeChatResp = JSON.parseObject(JSON.toJSONString(result.getData()), CozeChatResp.class);
            int maxCount = 20;
            int count = 0;
            String chatResult = null;
            while (count < maxCount) {
                boolean chatProgress = this.getChatProgress(cozeChatResp.getConversation_id(), cozeChatResp.getId(), secretToken);
                if (chatProgress) {
                    // 获取结果
                    chatResult = this.getChatResult(text, cozeChatResp.getConversation_id(), cozeChatResp.getId(), secretToken);
                    break;
                }
                count++;
                // 睡眠1秒
                Thread.sleep(1000);
            }
            // 需要简单判断一下是否合规，例如长度少于多少的时候或大于多少的时候则返回null
            log.info("coze text correction result:{}", chatResult);
            return this.formatResponse(chatResult, sourceText);
        } catch (Exception ex) {
            log.error("coze text correction result error", ex);
            log.error("coze text correction result error:{}", ex.getMessage());
            return null;
        } finally {
            CorrectUtils.printEndTime(config.getCode());
        }
    }

    private boolean getChatProgress(String conversationId, String chatId, String secretToken) {
        String url = String.format(GET_CHAT_DETAILS_URL, conversationId, chatId);
        try {
            // header中 增加bearer token
            HttpHeaders headers = new HttpHeaders();
            headers.setContentType(MediaType.APPLICATION_JSON);
            headers.set("Authorization", "Bearer " + secretToken);

            HttpEntity<String> entity = new HttpEntity<>(headers);

            ResponseEntity<CozeUnifiedResp> chatResp = this.restTemplate.exchange(url, HttpMethod.GET, entity, CozeUnifiedResp.class);
            CozeUnifiedResp body = chatResp.getBody();
            log.info("coze text correction progress result:{}", JSON.toJSONString(body));
            if (ObjectUtils.isEmpty(body) || body.getCode() != 0) {
                return false;
            }
            String data = JSON.toJSONString(body.getData());
            CozeChatResp chatProgressResp = JSON.parseObject(data, CozeChatResp.class);
            if (ObjectUtils.isEmpty(chatProgressResp) || !StringUtils.hasText(chatProgressResp.getStatus())) {
                return false;
            }
            log.info("coze对话进度:{}", chatProgressResp.getStatus());
            if (!chatProgressResp.getStatus().equals(CozeChatResp.STATUS_IN_PROGRESS)) {
                if (chatProgressResp.getStatus().equals(CozeChatResp.STATUS_COMPLETED)) {
                    return true;
                }
                // 如果解析失败，需要打印一下为什么失败了
                if (chatProgressResp.getStatus().equals(CozeChatResp.STATUS_FAILED)) {
                    String failChatResult = this.getChatResult(null, conversationId, chatId, secretToken);
                }
                log.warn("获取coze对话进度失败:{}", JSON.toJSONString(chatProgressResp));
                throw new BusinessException("获取coze对话进度失败:" + chatProgressResp.getStatus());
            }
            return false;
        } catch (Exception ex) {
            log.error("获取coze对话进度失败", ex);
            log.error("获取coze对话进度失败:{}", ex.getMessage());
            throw new BusinessException("获取coze对话进度失败:" + ex.getMessage());
        }
    }

    String getChatResult(String text, String conversationId, String chatId, String secretToken) {
        String url = String.format(GET_CHAT_RESULT_URL, conversationId, chatId);
        try {
            // header中 增加bearer token
            HttpHeaders headers = new HttpHeaders();
            headers.setContentType(MediaType.APPLICATION_JSON);
            headers.set("Authorization", "Bearer " + secretToken);

            HttpEntity<String> entity = new HttpEntity<>(headers);

            ResponseEntity<CozeUnifiedResp> chatResp = this.restTemplate.exchange(url, HttpMethod.GET, entity, CozeUnifiedResp.class);
            CozeUnifiedResp body = chatResp.getBody();
            if (ObjectUtils.isEmpty(body) || body.getCode() != 0) {
                return null;
            }
            String data = JSON.toJSONString(body.getData());
            if (!StringUtils.hasText(text)) {
                log.info("coze text correction fail result-v2:{}", data);
            }
            List<CozeChatResult> cozeChatResults = JSON.parseArray(data, CozeChatResult.class);
            log.info("coze chat result object:{}", JSON.toJSONString(cozeChatResults));
            if (ObjectUtils.isEmpty(cozeChatResults)) {
                return null;
            }
            for (CozeChatResult items : cozeChatResults) {
                if (items.getType().equals(CozeChatResult.TYPE_ANSWER)) {
                    // 如果他的答案超过了 源文本 + 5 个字，则直接返回，否则返回null
                    if (items.getContent().length() > text.length() + 5) {
                        return null;
                    }
                    return items.getContent();
                }
            }
            return null;
        } catch (Exception ex) {
            log.error("获取coze对话结果失败", ex);
            log.error("获取coze对话结果失败:{}", ex.getMessage());
            throw new BusinessException("获取coze对话结果失败:" + ex.getMessage());
        }
    }

    @Data
    static class CozeRequestParam {
        private String bot_id;
        private String user_id;
        private boolean stream;
        private boolean auto_save_history;
        List<additionalMessages> additional_messages = new ArrayList<>();

        public CozeRequestParam(String text, String botId) {
            this.bot_id = botId;
            this.user_id = "yunban";
            this.stream = false;
            this.auto_save_history = true;
            this.additional_messages.add(new additionalMessages(text));
        }
    }

    @Data
    static class additionalMessages {
        private String role;
        private String content;
        private String content_type;

        public additionalMessages(String content) {
            this.role = "user";
            this.content = content;
            this.content_type = "text";
        }
    }
}
